﻿using Common;
using Linguist.Acoustic;
using Linguist.G2p;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using Util;
using Util.Props;

namespace Linguist.Dictionary
{
    /**
    /// Creates a dictionary by quickly reading in an ASCII-based Sphinx-3 format
    /// dictionary. It is called the FastDictionary because the loading is fast. When
    /// loaded the dictionary just loads each line of the dictionary into the hash
    /// table, assuming that most words are not going to be used. Only when a word is
    /// actually used is its pronunciations massaged into an array of pronunciations.
    /// <p/>
    /// The format of the ASCII dictionary that it explains is the same as the
    /// {@link FullDictionary FullDictionary}, i.e., the word, followed by spaces or
    /// tab, followed by the pronunciation(s). For example, a digits dictionary will
    /// look like:
    /// <p/>
    /// 
    /// <pre>
    ///  ONE HH W AH N
    ///  ONE(2) W AH N
    ///  TWO T UW
    ///  THREE TH R IY
    ///  FOUR F AO R
    ///  FIVE F AY V
    ///  SIX S IH K S
    ///  SEVEN S EH V AH N
    ///  EIGHT EY T
    ///  NINE N AY N
    ///  ZERO Z IH R OW
    ///  ZERO(2) Z IY R OW
    ///  OH OW
    /// </pre>
    /// <p/>
    /// <p/>
    /// In the above example, the words "one" and "zero" have two pronunciations
    /// each.
     */
    public class FastDictionary: IDictionary
    {
        // -------------------------------
        // Configuration data
        // --------------------------------

        private Boolean addSilEndingPronunciation;
        private Boolean allowMissingWords;
        private Boolean createMissingWords;
        private String wordReplacement;
        protected String wordDictionaryFile;
        protected String fillerDictionaryFile;
        protected String g2pModelFile;
        protected int g2pMaxPron = 0;
        protected List<String> addendaUrlList;

        protected UnitManager unitManager;

        // -------------------------------
        // working data
        // -------------------------------
        protected Dictionary<String, String> dictionary;
        protected Dictionary<String, IWord> wordDictionary;
        protected G2PConverter g2pDecoder;

        protected static String FILLER_TAG = "-F-";
        protected List<String> fillerWords;
        protected Boolean allocated;


        public FastDictionary(String wordDictionaryFile, String fillerDictionaryFile,
                List<String> addendaUrlList, Boolean addSilEndingPronunciation,
                String wordReplacement, Boolean allowMissingWords,
                Boolean createMissingWords, UnitManager unitManager) 
        {
            this.wordDictionaryFile = wordDictionaryFile;
            this.fillerDictionaryFile = fillerDictionaryFile;
            this.addendaUrlList = addendaUrlList;
            this.addSilEndingPronunciation = addSilEndingPronunciation;
            this.wordReplacement = wordReplacement;
            this.allowMissingWords = allowMissingWords;
            this.createMissingWords = createMissingWords;
            this.unitManager = unitManager;
        }

        public FastDictionary(String wordDictionaryFile, String fillerDictionaryFile,
                List<String> addendaUrlList, Boolean addSilEndingPronunciation,
                String wordReplacement, Boolean allowMissingWords,
                Boolean createMissingWords, UnitManager unitManager,
                String g2pModelFile, int g2pMaxPron) 
        
            :this(wordDictionaryFile,
                 fillerDictionaryFile,
                 addendaUrlList,
                 addSilEndingPronunciation,
                 wordReplacement,
                 allowMissingWords,
                 createMissingWords,
                 unitManager)
        {
            this.g2pModelFile = g2pModelFile;
            this.g2pMaxPron = g2pMaxPron;
        }

        public FastDictionary() {

        }

        /*
        /// (non-Javadoc)
        /// 
        /// @see
        /// edu.cmu.sphinx.util.props.Configurable#newProperties(edu.cmu.sphinx.util
        /// .props.PropertySheet)
         */

        override public void newProperties(PropertySheet ps)
        {

            wordDictionaryFile = ConfigurationManagerUtils
                    .getResource(PROP_DICTIONARY, ps);
            fillerDictionaryFile = ConfigurationManagerUtils
                    .getResource(PROP_FILLER_DICTIONARY, ps);
            addendaUrlList = ps.getResourceStringList(PROP_ADDENDA);
            addSilEndingPronunciation = ps
                    .getBoolean(PROP_ADD_SIL_ENDING_PRONUNCIATION);
            wordReplacement = ps.getString(IDictionary.PROP_WORD_REPLACEMENT);
            allowMissingWords = ps.getBoolean(IDictionary.PROP_ALLOW_MISSING_WORDS);
            createMissingWords = ps.getBoolean(PROP_CREATE_MISSING_WORDS);
            unitManager = (UnitManager) ps.getComponent(PROP_UNIT_MANAGER);
            g2pModelFile = ConfigurationManagerUtils
                    .getResource(PROP_G2P_MODEL_PATH, ps);
            g2pMaxPron = ps.getInt(PROP_G2P_MAX_PRONUNCIATIONS);
        }

        /**
        /// Get the word dictionary file
        /// 
        /// @return the URL of the word dictionary file
         */
        public String getWordDictionaryFile() 
        {
            return wordDictionaryFile;
        }

        /**
        /// Get the filler dictionary file
        /// 
        /// @return the URL of the filler dictionary file
         */
        public String getFillerDictionaryFile() 
        {
            return fillerDictionaryFile;
        }

        /*
        /// (non-Javadoc)
        /// 
        /// @see edu.cmu.sphinx.linguist.dictionary.Dictionary#allocate()
         */

        override public void allocate()
        {
            if (!allocated) {
                dictionary = new Dictionary<String, String>();
                wordDictionary = new Dictionary<String, IWord>();

                Timer loadTimer = TimerPool.getTimer(this, "Load Dictionary");
                fillerWords = new List<String>();

                loadTimer.start();

                Trace.TraceInformation("Loading dictionary from: " + wordDictionaryFile);

                loadDictionary(new FileStream(wordDictionaryFile,FileMode.Open), false);

                loadCustomDictionaries(addendaUrlList);

                Trace.WriteLine("Loading filler dictionary from: "
                        + fillerDictionaryFile);

                loadDictionary(new FileStream(fillerDictionaryFile,FileMode.Open), true);

                if (g2pModelFile != null && !g2pModelFile.Equals("")) {
                    g2pDecoder = new G2PConverter(g2pModelFile);
                }
                loadTimer.stop();
            }

        }

        /*
        /// (non-Javadoc)
        /// 
        /// @see edu.cmu.sphinx.linguist.dictionary.Dictionary#deallocate()
         */

        override public void deallocate() 
        {
            if (allocated) {
                dictionary = null;
                g2pDecoder = null;
                allocated = false;
            }
        }

        /**
        /// Loads the given sphinx3 style simple dictionary from the given
        /// InputStream. The InputStream is assumed to contain ASCII data.
        /// 
        /// @param inputStream
        ///            the InputStream of the dictionary
        /// @param isFillerDict
        ///            true if this is a filler dictionary, false otherwise
        /// @throws java.io.IOException
        ///             if there is an error reading the dictionary
         */
        protected void loadDictionary(FileStream inputStream, Boolean isFillerDict)
        {
            
            StreamReader rd = new StreamReader(inputStream);
            //BufferedReader br = new BufferedReader(isr);
            String line;

            while ((line = rd.ReadLine()) != null) 
            {
                if (line.Length!=0) 
                {
                    int spaceIndex = line.IndexOf(' ');
                    int spaceIndexTab = line.IndexOf('\t');
                    if (spaceIndex == -1) 
                    {
                        // Case where there's no blank character
                        spaceIndex = spaceIndexTab;
                    } 
                    else if ((spaceIndexTab >= 0) && (spaceIndexTab < spaceIndex)) 
                    {
                        // Case where there's a blank and a tab, but the tab
                        // precedes the blank
                        spaceIndex = spaceIndexTab;
                    }
                    // TODO: throw an exception if spaceIndex == -1 ?
                    if (spaceIndex == -1) 
                    {
                        throw new Exception("Error loading word: " + line);
                    }
                    String word = line.Substring(0, spaceIndex);
                    word = word.ToLower();
                    // Add numeric index if the word is repeating.
                    if (dictionary.ContainsKey(word)) 
                    {
                        int index = 2;
                        String wordWithIdx;
                        do {
                            wordWithIdx = String.Format("%s(%d)", word, index++);
                        } 
                        while (dictionary.ContainsKey(wordWithIdx));
                        word = wordWithIdx;
                    }

                    if (isFillerDict) 
                    {
                        dictionary.Add(word, (FILLER_TAG + line));
                        fillerWords.Add(word);
                    } 
                    else 
                    {
                        dictionary.Add(word, line);
                    }
                }
            }

            rd.Close();
            inputStream.Close();
        }

        /**
        /// Gets a context independent unit. There should only be one instance of any
        /// CI unit
        /// 
        /// @param name
        ///            the name of the unit
        /// @param isFiller
        ///            if true, the unit is a filler unit
        /// @return the unit
         */
        protected Unit getCIUnit(String name, Boolean isFiller) 
        {
            return unitManager.getUnit(name, isFiller, Context.EMPTY_CONTEXT);
        }

        /**
        /// Returns the sentence start word.
        /// 
        /// @return the sentence start word
         */
        override public IWord getSentenceStartWord() 
        {
            return getWord(SENTENCE_START_SPELLING);
        }

        /**
        /// Returns the sentence end word.
        /// 
        /// @return the sentence end word
         */
        override public IWord getSentenceEndWord() 
        {
            return getWord(SENTENCE_END_SPELLING);
        }

        /**
        /// Returns the silence word.
        /// 
        /// @return the silence word
         */
        override public IWord getSilenceWord() 
        {
            return getWord(SILENCE_SPELLING);
        }

        /**
        /// Returns a Word object based on the spelling and its classification. The
        /// behavior of this method is also affected by the properties
        /// wordReplacement, allowMissingWords, and createMissingWords.
        /// 
        /// @param text
        ///            the spelling of the word of interest.
        /// @return a Word object
        /// @see edu.cmu.sphinx.linguist.dictionary.Word
         */
        override public IWord getWord(String text) 
        {
            text = text.ToLower();
            if(wordDictionary.ContainsKey(text))
            {
                return wordDictionary[text];
            }
            IWord wordObject = null;
        
            if (!dictionary.ContainsKey(text)) 
            { // deal with 'not found' case
                Trace.TraceInformation("The dictionary is missing a phonetic transcription for the word '"  + text + "'");
                if (wordReplacement != null) 
                {
                    wordObject = getWord(wordReplacement);
                } 
                else if (allowMissingWords) 
                {
                    if (createMissingWords) {
                        if (g2pModelFile != null
                                && !g2pModelFile.Equals("")) 
                        {
                            Trace.TraceInformation("Generating phonetic transcription(s) for the word '"
                                    + text + "' using g2p model");
                            List<G2p.Path> paths = g2pDecoder.phoneticize(text, g2pMaxPron);
                            List<Pronunciation> pronunciations = new List<Pronunciation>();
                            foreach (G2p.Path p in paths) 
                            {
                                int unitCount = p.getPath().Count;
                                List<IUnit> units = new List<IUnit>(unitCount);
                                foreach (String token in p.getPath()) 
                                {
                                    units.Add(getCIUnit(token, false));
                                }
                                if (units.Count == 0) 
                                {
                                    units.Add(UnitManager.SILENCE);
                                }
                                pronunciations.Add(new Pronunciation(units));
                                if (addSilEndingPronunciation) 
                                {
                                    units.Add(UnitManager.SILENCE);
                                    pronunciations.Add(new Pronunciation(units));
                                }
                            }
                            Pronunciation[] pronunciationsArray = pronunciations.ToArray();
                            wordObject = createWord(text,
                                                    pronunciationsArray,
                                                    false);
                            foreach (Pronunciation pronunciation in pronunciationsArray) 
                            {
                                pronunciation.setWord(wordObject);
                            }
                            wordDictionary.Add(text, wordObject);
                        } 
                        else 
                        {
                            wordObject = createWord(text, null, false);
                        }
                    }
                }
            } 
            else 
            { // first lookup for this string
                wordObject = processEntry(text);
            }

            return wordObject;
        }

        /**
        /// Create a Word object with the given spelling and pronunciations, and
        /// insert it into the dictionary.
        /// 
        /// @param text
        ///            the spelling of the word
        /// @param pronunciation
        ///            the pronunciation of the word
        /// @param isFiller
        ///            if <code>true</code> this is a filler word
        /// @return the word
         */
        private Word createWord(String text,
                                Pronunciation[] pronunciation,
                                Boolean isFiller) 
        {
            Word word = new Word(text, pronunciation, isFiller);
            if (!dictionary.ContainsKey(text))
                dictionary.Add(text, word.ToString());
            else
                dictionary[text] = word.ToString();
            return word;
        }

        /**
        /// Processes a dictionary entry. When loaded the dictionary just loads each
        /// line of the dictionary into the hash table, assuming that most words are
        /// not going to be used. Only when a word is actually used is its
        /// pronunciations massaged into an array of pronunciations.
         */
        private Word processEntry(String word) 
        {
            List<Pronunciation> pronunciations = new List<Pronunciation>();
            String line;
            int count = 0;
            Boolean isFiller = false;

            do {
                count++;
                String lookupWord = word;
                if (count > 1) {
                    lookupWord = lookupWord + '(' + count + ')';
                }

                if (dictionary.ContainsKey(lookupWord))
                {
                    line = dictionary[lookupWord];
                    String[] st = line.Split(new Char[] { ' ' });

                    String tag = st[0];
                    isFiller = tag.StartsWith(FILLER_TAG);
                    int unitCount = st.Length;

                    List<IUnit> units = new List<IUnit>(unitCount);
                    for (int i = 0; i < unitCount; i++)
                    {
                        String unitName = st[i];
                        units.Add(getCIUnit(unitName, isFiller));
                    }
                    pronunciations.Add(new Pronunciation(units));
                    if (!isFiller && addSilEndingPronunciation)
                    {
                        units.Add(UnitManager.SILENCE);
                        pronunciations.Add(new Pronunciation(units));
                    }
                }
                else
                    line = null;
            } while (line != null);

            Pronunciation[] pronunciationsArray = pronunciations.ToArray();
            Word wordObject = createWord(word, pronunciationsArray, isFiller);

            foreach (Pronunciation pronunciation in pronunciationsArray) 
            {
                pronunciation.setWord(wordObject);
            }
            wordDictionary.Add(word, wordObject);

            return wordObject;
        }

        /**
        /// Returns the set of all possible word classifications for this dictionary.
        /// 
        /// @return the set of all possible word classifications
         */
        override public IWordClassification[] getPossibleWordClassifications() 
        {
            return null;
        }

        /**
        /// Returns a string representation of this FastDictionary in alphabetical
        /// order.
        /// 
        /// @return a string representation of this FastDictionary
         */
        override
        public String ToString() 
        {

            SortedDictionary<String, String> sorted = new SortedDictionary<String, String>(dictionary);
            StringBuilder result = new StringBuilder();

            foreach (KeyValuePair<String, String> entry in sorted) 
            {
                result.Append(entry.Key);
                result.Append("   ").Append(entry.Value).Append('\n');
            }

            return result.ToString();
        }

        /**
        /// Gets the set of all filler words in the dictionary
        /// 
        /// @return an array (possibly empty) of all filler words
         */
        override public IWord[] getFillerWords() 
        {
            IWord[] fillerWordArray = new Word[fillerWords.Count];
            int index = 0;
            foreach (String spelling in fillerWords) 
            {
                fillerWordArray[index++] = getWord(spelling);
            }
            return fillerWordArray;
        }

        /**
        /// Dumps this FastDictionary to System.out.
         */
        public void dump() 
        {
            Trace.WriteLine(ToString());
        }

        /**
        /// Loads the dictionary with a list of URLs to custom dictionary resources
        /// 
        /// @param addenda
        ///            the list of custom dictionary URLs to be loaded
        /// @throws IOException
        ///             if there is an error reading the resource URL
         */
        private void loadCustomDictionaries(List<String> addenda)
        {
            if (addenda != null) 
            {
                foreach (String addendumUrl in addenda) 
                {
                    Trace.TraceInformation("Loading addendum dictionary from: " + addendumUrl);
                    loadDictionary(new FileStream(addendumUrl,FileMode.Open), false);
                }
            }
        }

    }
}
